### import ###
import xarray as xr
import numpy as np
from scipy.interpolate import griddata



### init ###
ear5path   = "./era5_2022060100_plev.nc"
f1 = xr.open_dataset(ear5path,engine="netcdf4")

zs = f1.zs[0].values/9.8  # 181 360  
lat  = f1.latitude[::-1]
lon = f1.longitude-180
ny = len(lat)
nx = len(lon)
lon_grid,lat_grid = np.meshgrid(lon,lat)

### interp

x = np.linspace(-180,180,21601)
y = np.linspace(-90,90,10801)
x_grid,y_grid = np.meshgrid(x,y)

points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))  # (700*1400, 2)
values = zs.ravel()  
grid_z = griddata(points, values, (x_grid, y_grid), method='cubic')


### inertp  mask

file_path = "./ETOPO1_Ice_g_gmt4.grd" 
data = xr.open_dataset(file_path)
mask = data["mask"].values  
lat = data["y"].values  
lon = data["x"].values  
lon_grid,lat_grid = np.meshgrid(lon,lat)

points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))  # (700*1400, 2)
values = mask.ravel()  

grid_mask = griddata(points, values,(x_grid,y_grid),method="nearest")





### make grd ###
coords = {'x': x}
lon = xr.DataArray(
    data=x,  
    dims=['x'],     
    coords=coords,  
    attrs={         
        'long_name': 'Longitude',
        'actual_range': [-180.0, 180.0],
        'units': 'degrees_east'
    }
)

coords = {'y': y}
lat = xr.DataArray(
    data=y,  
    dims=['y'],     
    coords=coords,  
    attrs={         
        'long_name': 'Latitude',
        'actual_range': [-90.0, 90.0],
        'units': 'degrees_north'
    }
)

coords = {'y':y,'x':x}
z = xr.DataArray(
    data=grid_z,  
    dims=['y','x'],     
    coords=coords, 
    attrs={         
        'long_name': 'z',
    }
)

coords = {'y':y,'x':x}
mask = xr.DataArray(
    data=grid_mask,  
    dims=['y','x'],     
    coords=coords,  
    attrs={         
        'long_name': 'Land mask (0: ocean, 1: land)',
    }
)


dataset = xr.Dataset(
    data_vars={
        'lon': lon,
        'lat': lat,
        'z': z,
        'mask':mask
    }
)

###  save grd  ###


dataset.to_netcdf("era5topo.grd")







